import sqlite3
from flask import current_app, g
from typing import Dict, List, Optional, Any
from datetime import datetime
import os
import json
import platform  # 新增平台模块导入
import ctypes    # 新增ctypes模块导入
from sqlalchemy.orm import sessionmaker


class DatabaseError(Exception):
    """数据库操作异常"""
    pass


def get_db():
    """获取数据库连接"""
    try:
        if 'db' not in g:
            g.db = sqlite3.connect(
                current_app.config['DATABASE'],
                detect_types=sqlite3.PARSE_DECLTYPES
            )
            g.db.row_factory = sqlite3.Row
        return g.db
    except Exception as e:
        current_app.logger.error(f"数据库连接错误: {str(e)}")
        raise DatabaseError(f"数据库连接失败: {str(e)}")

def close_db(e=None):
    """关闭数据库连接"""
    db = g.pop('db', None)
    if db is not None:
        try:
            db.close()
        except Exception as e:
            current_app.logger.error(f"关闭数据库连接错误: {str(e)}")

def init_db(app):
    """初始化数据库"""
    with app.app_context():
        conn = get_db()
        c = conn.cursor()
        
        # 创建下载表（添加status枚举约束）
        c.execute('''
            CREATE TABLE IF NOT EXISTS downloads (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                url TEXT NOT NULL,
                name TEXT,
                save_dir TEXT,
                status TEXT DEFAULT 'pending' CHECK(status IN (
                    'pending', 'preparing', 'downloading', 
                    'merging', 'completed', 'stopped',
                    'error', 'retrying', 'deleted'
                )),
                progress REAL DEFAULT 0,
                speed REAL DEFAULT 0,
                size INTEGER DEFAULT 0,
                completed INTEGER DEFAULT 0,
                total INTEGER DEFAULT 0,
                createdtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updatedtime TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                message TEXT,
                format TEXT,
                quality TEXT,
                download_mode TEXT,
                proxy TEXT,
                headers TEXT,
                cookies TEXT
            )
        ''')
        
        # 创建设置表
        c.execute('''
            CREATE TABLE IF NOT EXISTS settings (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL
            )
        ''')
        
        # 插入默认设置
        default_settings = [
            ('defaultSaveDirectory', 'downloads'),
            ('maxConcurrentDownloads', '3'),
            ('autoStartDownloads', 'true'),
            ('proxyEnabled', 'false'),
            ('proxyUrl', ''),
            ('connectionTimeout', '30'),
            ('retryCount', '3'),
            ('downloadMode', 'segmented'),
            ('defaultFormat', 'mp4'),
            ('defaultQuality', 'best'),
            ('customHeaders', '')
        ]
        
        c.executemany('INSERT OR IGNORE INTO settings (key, value) VALUES (?, ?)', default_settings)
        conn.commit()

def add_download(url: str, filename: str, save_dir: str) -> int:
    """添加下载任务"""
    try:
        db = get_db()
        cursor = db.execute(
            '''
            INSERT INTO downloads (
                url, filename, save_dir, status, created_at, updated_at
            ) VALUES (?, ?, ?, 'pending', CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
            ''',
            (url, filename, save_dir)
        )
        db.commit()
        return cursor.lastrowid
    except Exception as e:
        current_app.logger.error(f"添加下载任务错误: {str(e)}")
        raise DatabaseError(f"添加下载任务失败: {str(e)}")

def update_download(task_id: int, **kwargs) -> bool:
    """更新下载任务"""
    if not kwargs:
        return False
        
    try:
        # 添加状态转换验证
        if 'status' in kwargs:
            valid_statuses = {'pending', 'preparing', 'downloading', 'merging', 
                            'completed', 'stopped', 'error', 'retrying', 'deleted'}
            if kwargs['status'] not in valid_statuses:
                raise ValueError(f"无效状态: {kwargs['status']}")
        
        db = get_db()
        fields = []
        values = []
        
        for key, value in kwargs.items():
            fields.append(f"{key} = ?")
            values.append(value)
        
        values.append(task_id)
        query = f"UPDATE downloads SET {', '.join(fields)}, updated_at = CURRENT_TIMESTAMP WHERE id = ?"
        
        db.execute(query, values)
        db.commit()
        return True
    except Exception as e:
        current_app.logger.error(f"更新下载任务错误: {str(e)}")
        raise DatabaseError(f"更新下载任务失败: {str(e)}")

def get_download(task_id: int) -> Optional[Dict[str, Any]]:
    """获取下载任务"""
    try:
        db = get_db()
        row = db.execute('SELECT * FROM downloads WHERE id = ?', (task_id,)).fetchone()
        return dict(row) if row else None
    except Exception as e:
        current_app.logger.error(f"获取下载任务错误: {str(e)}")
        raise DatabaseError(f"获取下载任务失败: {str(e)}")

def get_all_downloads() -> List[Dict[str, Any]]:
    """获取所有下载任务"""
    try:
        db = get_db()
        rows = db.execute('SELECT * FROM downloads ORDER BY created_at DESC').fetchall()
        return [dict(row) for row in rows]
    except Exception as e:
        current_app.logger.error(f"获取所有下载任务错误: {str(e)}")
        raise DatabaseError(f"获取所有下载任务失败: {str(e)}")

def delete_download(task_id: int) -> bool:
    """删除下载任务"""
    try:
        db = get_db()
        db.execute('DELETE FROM downloads WHERE id = ?', (task_id,))
        db.commit()
        return True
    except Exception as e:
        current_app.logger.error(f"删除下载任务错误: {str(e)}")
        raise DatabaseError(f"删除下载任务失败: {str(e)}")

def get_setting(key: str) -> Optional[str]:
    """获取设置值"""
    try:
        db = get_db()
        row = db.execute('SELECT value FROM settings WHERE key = ?', (key,)).fetchone()
        return row['value'] if row else None
    except Exception as e:
        current_app.logger.error(f"获取设置值错误: {str(e)}")
        raise DatabaseError(f"获取设置值失败: {str(e)}")

def update_setting(key: str, value: str) -> bool:
    """更新设置值"""
    try:
        db = get_db()
        db.execute(
            'UPDATE settings SET value = ?, updated_at = CURRENT_TIMESTAMP WHERE key = ?',
            (value, key)
        )
        db.commit()
        return True
    except Exception as e:
        current_app.logger.error(f"更新设置值错误: {str(e)}")
        raise DatabaseError(f"更新设置值失败: {str(e)}")

def get_all_settings():
    """获取所有设置"""
    conn = get_db()
    c = conn.cursor()
    c.execute('SELECT key, value FROM settings')
    settings = dict(c.fetchall())
    
    # 将扁平结构转换为嵌套结构
    return {
        'basicSettings': {
            'defaultSaveDirectory': settings.get('defaultSaveDirectory', 'downloads'),
            'maxConcurrentDownloads': int(settings.get('maxConcurrentDownloads', '3')),
            'autoStartDownloads': settings.get('autoStartDownloads', 'true').lower() == 'true'
        },
        'networkSettings': {
            'proxyEnabled': settings.get('proxyEnabled', 'false').lower() == 'true',
            'proxyUrl': settings.get('proxyUrl', ''),
            'connectionTimeout': int(settings.get('connectionTimeout', '30')),
            'retryCount': int(settings.get('retryCount', '3'))
        },
        'advancedSettings': {
            'downloadMode': settings.get('downloadMode', 'segmented'),
            'defaultFormat': settings.get('defaultFormat', 'mp4'),
            'defaultQuality': settings.get('defaultQuality', 'best'),
            'customHeaders': settings.get('customHeaders', '')
        }
    }

def update_settings(settings_dict):
    """批量更新设置"""
    conn = get_db()
    c = conn.cursor()
    
    # 将嵌套结构转换为扁平结构
    flat_settings = {}
    for group, settings in settings_dict.items():
        for key, value in settings.items():
            flat_settings[key] = str(value)
    
    # 批量更新设置
    c.executemany('UPDATE settings SET value = ? WHERE key = ?',
                 [(value, key) for key, value in flat_settings.items()])
    conn.commit()

def save_download_state(task_id: int, state: Dict[str, Any]) -> None:
    """保存下载状态"""
    try:
        state_file = os.path.join(
            current_app.config['DOWNLOAD_DIR'],
            f'state_{task_id}.json'
        )
        with open(state_file, 'w') as f:
            json.dump(state, f)
    except Exception as e:
        current_app.logger.error(f"保存下载状态错误: {str(e)}")
        raise DatabaseError(f"保存下载状态失败: {str(e)}")

def load_download_state(task_id: int) -> Optional[Dict[str, Any]]:
    """加载下载状态"""
    try:
        state_file = os.path.join(
            current_app.config['DOWNLOAD_DIR'],
            f'state_{task_id}.json'
        )
        if os.path.exists(state_file):
            with open(state_file, 'r') as f:
                return json.load(f)
        return None
    except Exception as e:
        current_app.logger.error(f"加载下载状态错误: {str(e)}")
        raise DatabaseError(f"加载下载状态失败: {str(e)}")

def check_disk_space() -> bool:
    """检查磁盘空间（跨平台版本）"""
    try:
        if platform.system() == 'Windows':
            # Windows系统使用ctypes获取磁盘空间
            free_bytes = ctypes.c_ulonglong()
            total_bytes = ctypes.c_ulonglong()
            ctypes.windll.kernel32.GetDiskFreeSpaceExW(
                ctypes.c_wchar_p(current_app.config['DOWNLOAD_DIR']),
                ctypes.byref(free_bytes),
                ctypes.byref(total_bytes),
                None
            )
            free_space = free_bytes.value
        else:
            # Unix/Linux系统保持原有逻辑
            stat = os.statvfs(current_app.config['DOWNLOAD_DIR'])
            free_space = stat.f_frsize * stat.f_bavail
        
        min_space = int(get_setting('min_free_space') or 1073741824)
        return free_space >= min_space
    except Exception as e:
        current_app.logger.error(f"检查磁盘空间错误: {str(e)}")
        raise DatabaseError(f"检查磁盘空间失败: {str(e)}")

def update_download_progress(task_id: int, progress: float):
    """更新下载进度百分比"""
    try:
        db = get_db()
        db.execute(
            '''
            UPDATE downloads 
            SET progress = ?, updatedtime = CURRENT_TIMESTAMP 
            WHERE id = ?
            ''',
            (round(progress, 2), task_id)
        )
        db.commit()
    except Exception as e:
        current_app.logger.error(f"更新下载进度错误: {str(e)}")
        raise DatabaseError(f"更新进度失败: {str(e)}")

# Add this after get_db() function
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=get_db)